from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import datetime  # For datetime objects
import os.path  # To manage paths
import sys  # To find out the script name (in argv[0])

# 获取当前脚本的绝对路径，并找到父目录
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
# 将父目录添加到 Python 模块搜索路径
sys.path.append(parent_dir)
# Import the backtrader platform
import backtrader as bt

# 创建策略继承bt.Strategy 
class TestStrategy(bt.Strategy): 
    def log(self, txt, dt=None): 
        # 记录策略的执行日志  
        dt = dt or self.datas[0].datetime.date(0) 
        print('%s, %s' % (dt.isoformat(), txt))

    def __init__(self): 
    # 保存收盘价的引用  
        self.dataclose = self.datas[0].close 
        self.dataopen = self.datas[0].open
        # 跟踪挂单  
        self.order = None 

    # buy()和sell()会返回被创建的订单(尚未执行的)， 订单状态的更改将通过 notify方法通知给策略Strategy。  
    def notify_order(self, order): 
        if order.status in [order.Submitted, order.Accepted]: 
            # broker 提交/接受了，买/卖订单则什么都不做  
            return
        # 检查一个订单是否完成  
        # 注意: 当资金不足时，broker会拒绝订单  
        if order.status in [order.Completed]: 
            if order.isbuy(): 
                self.log('已买入, %.2f' % order.executed.price) 
            elif order.issell(): 
                self.log('已卖出, %.2f' % order.executed.price) 

            # 记录买入时的K线位置
            self.bar_executed = len(self) 

        elif order.status in [order.Canceled, order.Margin, order.Rejected]: 
            self.log('订单取消/保证金不足/拒绝') 
            
        # 其他状态记录为：无挂起订单  
        self.order = None 


    def next(self): 
        # 记录收盘价  
        self.log('Open, %.2f, Close, %.2f' % (self.dataopen[0], self.dataclose[0])) 

        # 如果有订单正在挂起，不操作  
        if self.order: 
            return 

        # 如果没有持仓则买入 -  Strategy类有一个变量position保存当前持有的资产数量
        if not self.position: 
            # 今天的收盘价 < 昨天收盘价  
            if self.dataclose[0] < self.dataclose[-1]: 
                # 昨天收盘价 < 前天的收盘价  
                if self.dataclose[-1] < self.dataclose[-2]: 
                    # 买入  
                    self.log('买入, %.2f' % self.dataclose[0]) 
                    # 跟踪订单避免重复. 连续3天的收盘价是跌，则挂单买入（已经停市，所以实际用第四天的开盘价成交），买入单位是1, 价格是第四天的开盘价
                    self.order = self.buy() 
        else: 
            # 5个K线柱后（第6个K线柱）不管涨跌都卖。 请注意，这里没有指定具体时间，而是指定的柱的数量。
            # 一个柱可能代表 1分钟、1小时、1天、1星期等等，这取决于你价格数据文件里一条数据代表的周期。 虽然我们知道每个柱代表一天，但策略并不知道。
            # 如果已经持仓，len(self): 当前已处理的 K 线总数（从 0 开始计数）. 判断当前 K 线数量是否超过买入时的 K 线位置加 5, 不管涨跌都卖。  
            if len(self) >= (self.bar_executed + 5): 
                # 全部卖出  
                self.log('卖出, %.2f' % self.dataclose[0]) 
                # 跟踪订单避免重复. 卖出单位是1，第6个K柱的close的时候挂单，实际成交的价格是第7个K柱的开盘价
                self.order = self.sell() 

if __name__ == '__main__':
    # 创建Cerebro引擎  
    cerebro = bt.Cerebro() 
    # Cerebro引擎在后台创建broker(经纪人)，系统默认资金量为10000 
    # 为Cerebro引擎添加策略  
    cerebro.addstrategy(TestStrategy) 
    # 获取当前运行脚本所在目录  
    modpath = os.path.dirname(os.path.abspath(sys.argv[0])) 
    # 拼接加载路径  
    datapath = os.path.join(modpath, '../datas/orcl-1995-2014.txt') 
    # 创建交易数据集  
    data = bt.feeds.YahooFinanceCSVData( 
                                    dataname=datapath, 
                                    # 数据必须大于fromdate 
                                    fromdate=datetime.datetime(2000, 1, 1), 
                                    # 数据必须小于todate 
                                    todate=datetime.datetime(2000, 12, 31), 
                                    reverse=False) 
    # 加载交易数据  
    cerebro.adddata(data) 
    # 设置投资金额100000.0 
    cerebro.broker.setcash(100000.0) 
    # 引擎运行前打印期出资金  
    print('组合期初资金: %.2f' % cerebro.broker.getvalue()) 
    cerebro.run() 
    # 引擎运行后打印期末资金  
    print('组合期末资金: %.2f' % cerebro.broker.getvalue())